{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b55233d4",
   "metadata": {},
   "source": [
    "# Automatic Differentiation with `BrainPyObject`\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/brainpy/blob/master/docs_version2/tutorial_advanced/differentiation.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/brainpy/blob/master/docs_version2/tutorial_advanced/differentiation.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "355bb9b6",
   "metadata": {},
   "source": [
    "@[Chaoming Wang](https://github.com/chaoming0625)\n",
    "@[Xiaoyu Chen](mailto:c-xy17@tsinghua.org.cn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbc1e2d7",
   "metadata": {},
   "source": [
    "In this section, we are going to talk about how to realize automatic differentiation on your variables in a function or a class object. In current machine learning systems, gradients are commonly used in various situations. Therefore, we should understand:\n",
    "\n",
    "- How to calculate derivatives of arbitrary complex functions?\n",
    "- How to compute high-order gradients?"
   ]
  },
  {
   "cell_type": "code",
   "id": "e0ae6076",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.241942Z",
     "start_time": "2025-10-06T03:32:42.280993Z"
    }
   },
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math as bm\n",
    "\n",
    "# bm.set_platform('cpu')\n",
    "\n",
    "bp.__version__"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3.0.0'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "id": "7afa7421",
   "metadata": {},
   "source": [
    "## Preliminary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01ca8416",
   "metadata": {},
   "source": [
    "Every autograd function in BrainPy has several keywords. All examples below are illustrated through [brainpy.math.grad()](../apis/auto/math/generated/brainpy.math.autograd.grad.rst). Other autograd functions have the same settings. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d75313f",
   "metadata": {},
   "source": [
    "### ``argnums`` and ``grad_vars``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "772965c3",
   "metadata": {},
   "source": [
    "The autograd functions in BrainPy can compute derivatives of *function arguments* (specified by `argnums`) or *non-argument variables* (specified by ``grad_vars``). For instance, the following is a linear readout model:"
   ]
  },
  {
   "cell_type": "code",
   "id": "be17f596",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.450685Z",
     "start_time": "2025-10-06T03:32:46.250537Z"
    }
   },
   "source": [
    "class Linear(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(Linear, self).__init__()\n",
    "        self.w = bm.Variable(bm.random.random((1, 10)))\n",
    "        self.b = bm.Variable(bm.zeros(1))\n",
    "    \n",
    "    def update(self, x):\n",
    "        r = bm.dot(self.w, x) + self.b\n",
    "        return r.sum()\n",
    "    \n",
    "l = Linear()"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "id": "47d93392",
   "metadata": {},
   "source": [
    "If we try to focus on the derivative of the argument \"x\" when calling the update function, we can set this through ``argnums``:"
   ]
  },
  {
   "cell_type": "code",
   "id": "7bf6ae1f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.549403Z",
     "start_time": "2025-10-06T03:32:46.455785Z"
    }
   },
   "source": [
    "grad = bm.grad(l.update, argnums=0)\n",
    "\n",
    "grad(bm.ones(10))"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([0.9260962 , 0.09364974, 0.00942123, 0.9127095 , 0.67102027,\n",
       "                   0.620348  , 0.5872253 , 0.43891704, 0.7220038 , 0.9008471 ],      dtype=float32),\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "id": "7beb97b5",
   "metadata": {},
   "source": [
    "By contrast, if you focus on the derivatives of parameters \"self.w\" and \"self.b\", we should label them with ``grad_vars``:  "
   ]
  },
  {
   "cell_type": "code",
   "id": "f1f0d2c7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.593515Z",
     "start_time": "2025-10-06T03:32:46.559123Z"
    }
   },
   "source": [
    "grad = bm.grad(l.update, grad_vars=(l.w, l.b))\n",
    "\n",
    "grad(bm.ones(10))"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]], dtype=float32),\n",
       " Array([1.], dtype=float32))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "id": "17ea78df",
   "metadata": {},
   "source": [
    "If we pay attention to the derivatives of both argument \"x\" and parameters \"self.w\" and \"self.b\", ``argnums`` and ``grad_vars`` can be used together. In this condition, the gradient function will return gradients with the format of ``(var_grads, arg_grads)``, where ``arg_grads`` refers to the gradients of \"argnums\" and ``var_grads`` refers to the gradients of \"grad_vars\". "
   ]
  },
  {
   "cell_type": "code",
   "id": "5cc0347c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.606975Z",
     "start_time": "2025-10-06T03:32:46.597476Z"
    }
   },
   "source": [
    "grad = bm.grad(l.update, grad_vars=(l.w, l.b), argnums=0)\n",
    "\n",
    "var_grads, arg_grads = grad(bm.ones(10))"
   ],
   "outputs": [],
   "execution_count": 5
  },
  {
   "cell_type": "code",
   "id": "ce6f0f99",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.627520Z",
     "start_time": "2025-10-06T03:32:46.622567Z"
    }
   },
   "source": [
    "var_grads"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]], dtype=float32),\n",
       " Array([1.], dtype=float32))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "cell_type": "code",
   "id": "aa0d8b7f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.650348Z",
     "start_time": "2025-10-06T03:32:46.642056Z"
    }
   },
   "source": [
    "arg_grads"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([0.9260962 , 0.09364974, 0.00942123, 0.9127095 , 0.67102027,\n",
       "                   0.620348  , 0.5872253 , 0.43891704, 0.7220038 , 0.9008471 ],      dtype=float32),\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 7
  },
  {
   "cell_type": "markdown",
   "id": "93f20772",
   "metadata": {},
   "source": [
    "### ``return_value``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42e5b9dd",
   "metadata": {},
   "source": [
    "As is mentioned above, autograd functions return a function which computes gradients regardless of the returned value. Sometimes, however, we care about the value the function returns, not just the gradients. In this condition, you can set ``return_value=True`` in the autograd function."
   ]
  },
  {
   "cell_type": "code",
   "id": "600ea97e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.677529Z",
     "start_time": "2025-10-06T03:32:46.668211Z"
    }
   },
   "source": [
    "grad = bm.grad(l.update, argnums=0, return_value=True)\n",
    "\n",
    "gradient, value = grad(bm.ones(10))"
   ],
   "outputs": [],
   "execution_count": 8
  },
  {
   "cell_type": "code",
   "id": "d6909c04",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.700208Z",
     "start_time": "2025-10-06T03:32:46.694602Z"
    }
   },
   "source": [
    "gradient"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([0.9260962 , 0.09364974, 0.00942123, 0.9127095 , 0.67102027,\n",
       "                   0.620348  , 0.5872253 , 0.43891704, 0.7220038 , 0.9008471 ],      dtype=float32),\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 9
  },
  {
   "cell_type": "code",
   "id": "528b392f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.720441Z",
     "start_time": "2025-10-06T03:32:46.715681Z"
    }
   },
   "source": [
    "value"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(5.8822384, dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "id": "a5f829bd",
   "metadata": {},
   "source": [
    "### ``has_aux``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d9f4e2b",
   "metadata": {},
   "source": [
    "In some situations, we are interested in the intermediate values in a function, and ``has_aux=True`` can be of great help. The constraint is that you must return values with the format of ``(loss, aux_data)``. For instance, "
   ]
  },
  {
   "cell_type": "code",
   "id": "28e93b87",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.745938Z",
     "start_time": "2025-10-06T03:32:46.740676Z"
    }
   },
   "source": [
    "class LinearAux(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(LinearAux, self).__init__()\n",
    "        self.w = bm.Variable(bm.random.random((1, 10)))\n",
    "        self.b = bm.Variable(bm.zeros(1))\n",
    "    \n",
    "    def update(self, x):\n",
    "        dot = bm.dot(self.w, x)\n",
    "        r = (dot + self.b).sum()\n",
    "        return r, (r, dot)  # here the aux data is a tuple, includes the loss and the dot value.\n",
    "                            # however, aux can be arbitrary complex.\n",
    "    \n",
    "l2 = LinearAux()"
   ],
   "outputs": [],
   "execution_count": 11
  },
  {
   "cell_type": "code",
   "id": "3c683624",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.771383Z",
     "start_time": "2025-10-06T03:32:46.764029Z"
    }
   },
   "source": [
    "grad = bm.grad(l2.update, argnums=0, has_aux=True)\n",
    "\n",
    "gradient, aux = grad(bm.ones(10))"
   ],
   "outputs": [],
   "execution_count": 12
  },
  {
   "cell_type": "code",
   "id": "828ae73f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.781409Z",
     "start_time": "2025-10-06T03:32:46.777134Z"
    }
   },
   "source": [
    "gradient"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([0.61356115, 0.29959154, 0.5242523 , 0.5245906 , 0.42144704,\n",
       "                   0.5005964 , 0.42047822, 0.9298103 , 0.02631509, 0.35181725],      dtype=float32),\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 13
  },
  {
   "cell_type": "code",
   "id": "d921e0d6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.802647Z",
     "start_time": "2025-10-06T03:32:46.797560Z"
    }
   },
   "source": [
    "aux"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array(4.6124597, dtype=float32),\n",
       " Array(value=Array([4.6124597]), dtype=float32))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 14
  },
  {
   "cell_type": "markdown",
   "id": "6becdd17",
   "metadata": {},
   "source": [
    "When multiple keywords (``argnums``, ``grad_vars``, ``has_aux`` or``return_value``) are set simulatenously, the return format of the gradient function can be inspected through the corresponding API documentation [brainpy.math.grad()](../apis/auto/math/generated/brainpy.math.autograd.grad.rst)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df6b31f4",
   "metadata": {},
   "source": [
    "## ``brainpy.math.grad()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d289c868",
   "metadata": {},
   "source": [
    "[brainpy.math.grad()](../apis/auto/math/generated/brainpy.math.autograd.grad.rst) takes a function/object ($f : \\mathbb{R}^n \\to \\mathbb{R}$) as the input and returns a new function ($\\partial f(x) \\to \\mathbb{R}^n$) which computes the gradient of the original function/object. It's worthy to note that ``brainpy.math.grad()`` only supports returning scalar values. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56075f51",
   "metadata": {},
   "source": [
    "### Pure functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b4cccc",
   "metadata": {},
   "source": [
    "For pure function, the gradient is taken with respect to the first argument: "
   ]
  },
  {
   "cell_type": "code",
   "id": "45352485",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.823304Z",
     "start_time": "2025-10-06T03:32:46.817440Z"
    }
   },
   "source": [
    "def f(a, b):\n",
    "    return a * 2 + b\n",
    "\n",
    "grad_f1 = bm.grad(f)"
   ],
   "outputs": [],
   "execution_count": 15
  },
  {
   "cell_type": "code",
   "id": "6009405f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.886584Z",
     "start_time": "2025-10-06T03:32:46.840160Z"
    }
   },
   "source": [
    "grad_f1(2., 1.)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(2., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 16
  },
  {
   "cell_type": "markdown",
   "id": "c06f4f4e",
   "metadata": {},
   "source": [
    "However, this can be controlled via the `argnums` argument."
   ]
  },
  {
   "cell_type": "code",
   "id": "58aa6fbc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.911717Z",
     "start_time": "2025-10-06T03:32:46.902988Z"
    }
   },
   "source": [
    "grad_f2 = bm.grad(f, argnums=(0, 1))\n",
    "\n",
    "grad_f2(2., 1.)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array(2., dtype=float32, weak_type=True),\n",
       " Array(1., dtype=float32, weak_type=True))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 17
  },
  {
   "cell_type": "markdown",
   "id": "2f0874ef",
   "metadata": {},
   "source": [
    "### Class objects"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00906f22",
   "metadata": {},
   "source": [
    "For a class object or a class bound function, the gradient is taken with respect to the provided ``grad_vars`` and ``argnums`` setting: "
   ]
  },
  {
   "cell_type": "code",
   "id": "acc95d4c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:46.936904Z",
     "start_time": "2025-10-06T03:32:46.927317Z"
    }
   },
   "source": [
    "class F(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(F, self).__init__()\n",
    "        self.a = bm.TrainVar(bm.ones(1))\n",
    "        self.b = bm.TrainVar(bm.ones(1))\n",
    "\n",
    "    def __call__(self, c):\n",
    "        ab = self.a * self.b\n",
    "        ab2 = ab * 2\n",
    "        vv = ab2 + c\n",
    "        return vv.mean()\n",
    "    \n",
    "f = F()"
   ],
   "outputs": [],
   "execution_count": 18
  },
  {
   "cell_type": "markdown",
   "id": "18d64bc3",
   "metadata": {},
   "source": [
    "The ``grad_vars`` can be a Array, or a list/tuple/dict of Array. "
   ]
  },
  {
   "cell_type": "code",
   "id": "30484eab",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.041333Z",
     "start_time": "2025-10-06T03:32:46.953417Z"
    }
   },
   "source": [
    "bm.grad(f, grad_vars=f.train_vars())(10.)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'F0.a': Array([2.], dtype=float32), 'F0.b': Array([2.], dtype=float32)}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 19
  },
  {
   "cell_type": "code",
   "id": "fa99d3ef",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.070892Z",
     "start_time": "2025-10-06T03:32:47.061457Z"
    }
   },
   "source": [
    "bm.grad(f, grad_vars=[f.a, f.b])(10.)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Array([2.], dtype=float32), Array([2.], dtype=float32)]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 20
  },
  {
   "cell_type": "markdown",
   "id": "0847c77f",
   "metadata": {},
   "source": [
    "If there are dynamically changed values in the gradient function, you can provide them in the ``dyn_vars`` argument. "
   ]
  },
  {
   "cell_type": "code",
   "id": "f77b4c0e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.091003Z",
     "start_time": "2025-10-06T03:32:47.087472Z"
    }
   },
   "source": [
    "class F2(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(F2, self).__init__()\n",
    "        self.a = bm.TrainVar(bm.ones(1))\n",
    "        self.b = bm.TrainVar(bm.ones(1))\n",
    "\n",
    "    def __call__(self, c):\n",
    "        ab = self.a * self.b\n",
    "        ab = ab * 2\n",
    "        self.a.value = ab\n",
    "        return (ab + c).mean()"
   ],
   "outputs": [],
   "execution_count": 21
  },
  {
   "cell_type": "code",
   "id": "a0cf62b0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.121482Z",
     "start_time": "2025-10-06T03:32:47.108082Z"
    }
   },
   "source": [
    "f2 = F2()\n",
    "bm.grad(f2, grad_vars=f2.b)(10.)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([2.], dtype=float32)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 22
  },
  {
   "cell_type": "markdown",
   "id": "6998ec7c",
   "metadata": {},
   "source": [
    "Besides, if you are interested in the gradient of the input value, please use the ``argnums`` argument. Then, the gradient function will return ``(grads_of_grad_vars, grads_of_args)``. "
   ]
  },
  {
   "cell_type": "code",
   "id": "42c0dca2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.155859Z",
     "start_time": "2025-10-06T03:32:47.151726Z"
    }
   },
   "source": [
    "class F3(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(F3, self).__init__()\n",
    "        self.a = bm.TrainVar(bm.ones(1))\n",
    "        self.b = bm.TrainVar(bm.ones(1))\n",
    "\n",
    "    def __call__(self, c, d):\n",
    "        ab = self.a * self.b\n",
    "        ab = ab * 2\n",
    "        return (ab + c * d).mean()"
   ],
   "outputs": [],
   "execution_count": 23
  },
  {
   "cell_type": "code",
   "id": "3fe1c9ce",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.195751Z",
     "start_time": "2025-10-06T03:32:47.170453Z"
    }
   },
   "source": [
    "f3 = F3()\n",
    "grads_of_gv, grad_of_args = bm.grad(f3, grad_vars=[f3.a, f3.b], argnums=0)(10., 3.)\n",
    "\n",
    "print(\"grads_of_gv :\", grads_of_gv)\n",
    "print(\"grad_of_args :\", grad_of_args)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "grads_of_gv : [Array([2.], dtype=float32), Array([2.], dtype=float32)]\n",
      "grad_of_args : 3.0\n"
     ]
    }
   ],
   "execution_count": 24
  },
  {
   "cell_type": "code",
   "id": "ba55cac6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.216353Z",
     "start_time": "2025-10-06T03:32:47.202101Z"
    }
   },
   "source": [
    "f3 = F3()\n",
    "grads_of_gv, grad_of_args = bm.grad(f3, grad_vars=[f3.a, f3.b], argnums=(0, 1))(10., 3.)\n",
    "\n",
    "print(\"grads_of_gv :\", grads_of_gv)\n",
    "print(\"grad_of_args :\", grad_of_args)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "grads_of_gv : [Array([2.], dtype=float32), Array([2.], dtype=float32)]\n",
      "grad_of_args : (Array(3., dtype=float32, weak_type=True), Array(10., dtype=float32, weak_type=True))\n"
     ]
    }
   ],
   "execution_count": 25
  },
  {
   "cell_type": "markdown",
   "id": "06491457",
   "metadata": {},
   "source": [
    "Actually, it is recommended to provide all dynamically changed variables, whether or not they are updated in the gradient function, in the ``dyn_vars`` argument. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73cedb4d",
   "metadata": {},
   "source": [
    "### Auxiliary data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1469a67f",
   "metadata": {},
   "source": [
    "Usually, we want to get the loss value, or we want to return some intermediate variables during the gradient computation. In these situation, users can set ``has_aux=True`` to return auxiliary data and set ``return_value=True`` to return the loss value. "
   ]
  },
  {
   "cell_type": "code",
   "id": "a34a7e5a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.243950Z",
     "start_time": "2025-10-06T03:32:47.233092Z"
    }
   },
   "source": [
    "# return loss\n",
    "\n",
    "grad, loss = bm.grad(f, grad_vars=f.a, return_value=True)(10.)\n",
    "\n",
    "print('grad: ', grad)\n",
    "print('loss: ', loss)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "grad:  [2.]\n",
      "loss:  12.0\n"
     ]
    }
   ],
   "execution_count": 26
  },
  {
   "cell_type": "code",
   "id": "4a1ad862",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.277225Z",
     "start_time": "2025-10-06T03:32:47.266418Z"
    }
   },
   "source": [
    "class F4(bp.BrainPyObject):\n",
    "    def __init__(self):\n",
    "        super(F4, self).__init__()\n",
    "        self.a = bm.TrainVar(bm.ones(1))\n",
    "        self.b = bm.TrainVar(bm.ones(1))\n",
    "\n",
    "    def __call__(self, c):\n",
    "        ab = self.a * self.b\n",
    "        ab2 = ab * 2\n",
    "        loss = (ab + c).mean()\n",
    "        return loss, (ab, ab2)\n",
    "    \n",
    "\n",
    "f4 = F4()\n",
    "    \n",
    "# return intermediate values\n",
    "grad, aux_data = bm.grad(f4, grad_vars=f4.a, has_aux=True)(10.)\n",
    "\n",
    "print('grad: ', grad)\n",
    "print('aux_data: ', aux_data)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "grad:  [1.]\n",
      "aux_data:  (Array([1.], dtype=float32), Array([2.], dtype=float32))\n"
     ]
    }
   ],
   "execution_count": 27
  },
  {
   "cell_type": "markdown",
   "id": "33d2c322",
   "metadata": {},
   "source": [
    "```note\n",
    "Any function used to compute gradients through ``brainpy.math.grad()`` must return a scalar value. Otherwise an error will raise. \n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "id": "ea6a89f5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.340347Z",
     "start_time": "2025-10-06T03:32:47.301245Z"
    }
   },
   "source": [
    "try:\n",
    "    bm.grad(lambda x: x)(bm.zeros(2))\n",
    "except Exception as e:\n",
    "    print(type(e), e)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'TypeError'> Gradient only defined for scalar-output functions. Output had shape: (2,).\n"
     ]
    }
   ],
   "execution_count": 28
  },
  {
   "cell_type": "code",
   "id": "d08e3753",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.595314Z",
     "start_time": "2025-10-06T03:32:47.362810Z"
    }
   },
   "source": [
    "# this is right\n",
    "\n",
    "bm.grad(lambda x: x.mean())(bm.zeros(2))"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([0.5, 0.5]), dtype=float32)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 29
  },
  {
   "cell_type": "markdown",
   "id": "119967c0",
   "metadata": {},
   "source": [
    "## ``brainpy.math.vector_grad()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1542356e",
   "metadata": {},
   "source": [
    "If users want to take gradients for a vector-output values, please use the [brainpy.math.vector_grad()](../apis/auto/math/generated/brainpy.math.autograd.vector_grad.rst) function. For example, "
   ]
  },
  {
   "cell_type": "code",
   "id": "9a0a9b71",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.614576Z",
     "start_time": "2025-10-06T03:32:47.611323Z"
    }
   },
   "source": [
    "def f(a, b): \n",
    "    return bm.sin(b) * a"
   ],
   "outputs": [],
   "execution_count": 30
  },
  {
   "cell_type": "markdown",
   "id": "fcb68361",
   "metadata": {},
   "source": [
    "#### Gradients for vectors"
   ]
  },
  {
   "cell_type": "code",
   "id": "1323e89d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.714816Z",
     "start_time": "2025-10-06T03:32:47.621145Z"
    }
   },
   "source": [
    "# vectors\n",
    "\n",
    "a = bm.arange(5.)\n",
    "b = bm.random.random(5)"
   ],
   "outputs": [],
   "execution_count": 31
  },
  {
   "cell_type": "code",
   "id": "a776e614",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.771624Z",
     "start_time": "2025-10-06T03:32:47.720300Z"
    }
   },
   "source": [
    "bm.vector_grad(f)(a, b)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([0.31165043, 0.16343522, 0.28338438, 0.7613114 , 0.71029633],      dtype=float32)"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 32
  },
  {
   "cell_type": "code",
   "id": "85748195",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.809889Z",
     "start_time": "2025-10-06T03:32:47.786633Z"
    }
   },
   "source": [
    "bm.vector_grad(f, argnums=(0, 1))(a, b)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([0.31165043, 0.16343522, 0.28338438, 0.7613114 , 0.71029633],      dtype=float32),\n",
       " Array([0.       , 0.9865541, 1.9180129, 1.9451592, 2.8156111], dtype=float32))"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 33
  },
  {
   "cell_type": "markdown",
   "id": "10694945",
   "metadata": {},
   "source": [
    "#### Gradients for matrices"
   ]
  },
  {
   "cell_type": "code",
   "id": "19acd682",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:47.956877Z",
     "start_time": "2025-10-06T03:32:47.830058Z"
    }
   },
   "source": [
    "# matrix\n",
    "\n",
    "a = bm.arange(6.).reshape((2, 3))\n",
    "b = bm.random.random((2, 3))"
   ],
   "outputs": [],
   "execution_count": 34
  },
  {
   "cell_type": "code",
   "id": "4c049c25",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.060530Z",
     "start_time": "2025-10-06T03:32:47.980931Z"
    }
   },
   "source": [
    "bm.vector_grad(f, argnums=1)(a, b)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[0.       , 0.8268384, 1.4004612],\n",
       "       [2.697572 , 3.6508353, 4.748515 ]], dtype=float32)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 35
  },
  {
   "cell_type": "code",
   "id": "060fb4f2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.081092Z",
     "start_time": "2025-10-06T03:32:48.072443Z"
    }
   },
   "source": [
    "bm.vector_grad(f, argnums=(0, 1))(a, b)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([[0.29724902, 0.56243956, 0.7139168 ],\n",
       "        [0.43755704, 0.40861055, 0.31315213]], dtype=float32),\n",
       " Array([[0.       , 0.8268384, 1.4004612],\n",
       "        [2.697572 , 3.6508353, 4.748515 ]], dtype=float32))"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 36
  },
  {
   "cell_type": "markdown",
   "id": "55e96324",
   "metadata": {},
   "source": [
    "Similar to [brainpy.math.grad()](../apis/auto/math/generated/brainpy.math.autograd.grad.rst) , ``brainpy.math.vector_grad()`` also supports derivatives of variables in a class object. Here is a simple example. "
   ]
  },
  {
   "cell_type": "code",
   "id": "34e4f7bd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.103900Z",
     "start_time": "2025-10-06T03:32:48.098973Z"
    }
   },
   "source": [
    "class Test(bp.BrainPyObject):\n",
    "  def __init__(self):\n",
    "    super(Test, self).__init__()\n",
    "    self.x = bm.Variable(bm.ones(5))\n",
    "    self.y = bm.Variable(bm.ones(5))\n",
    "\n",
    "  def __call__(self):\n",
    "    return self.x ** 2 + self.y ** 3 + 10\n",
    "\n",
    "t = Test()"
   ],
   "outputs": [],
   "execution_count": 37
  },
  {
   "cell_type": "code",
   "id": "91fb638c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.265004Z",
     "start_time": "2025-10-06T03:32:48.123612Z"
    }
   },
   "source": [
    "bm.vector_grad(t, grad_vars=t.x)()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([2., 2., 2., 2., 2.], dtype=float32)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 38
  },
  {
   "cell_type": "code",
   "id": "678e2a24",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.307732Z",
     "start_time": "2025-10-06T03:32:48.296920Z"
    }
   },
   "source": [
    "bm.vector_grad(t, grad_vars=(t.x, ))()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([2., 2., 2., 2., 2.], dtype=float32),)"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 39
  },
  {
   "cell_type": "code",
   "id": "d3279ad8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.337759Z",
     "start_time": "2025-10-06T03:32:48.325688Z"
    }
   },
   "source": [
    "bm.vector_grad(t, grad_vars=(t.x, t.y))()"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array([2., 2., 2., 2., 2.], dtype=float32),\n",
       " Array([3., 3., 3., 3., 3.], dtype=float32))"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 40
  },
  {
   "cell_type": "markdown",
   "id": "25b9cb39",
   "metadata": {},
   "source": [
    "Other operations like ``return_value`` and ``has_aux`` in [brainpy.math.vector_grad()](../apis/auto/math/generated/brainpy.math.autograd.vector_grad.rst)  are the same as those in [brainpy.math.grad()](../apis/auto/math/generated/brainpy.math.autograd.grad.rst) ."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ca257d2",
   "metadata": {},
   "source": [
    "## ``brainpy.math.jacobian()``"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f68747a3",
   "metadata": {},
   "source": [
    "Another way to take gradients of a vector-output value is using [brainpy.math.jacobian()](../apis/auto/math/generated/brainpy.math.autograd.jacobian.rst). ``brainpy.math.jacobian()`` aims to automatically compute the Jacobian matrices $\\partial f(x) \\in \\mathbb{R}^{m \\times n}$ by the given function $f : \\mathbb{R}^n \\to \\mathbb{R}^m$ at the given point of $x \\in \\mathbb{R}^n$. Here, we will not go to the details of the implementation and usage of the ``brainpy.math.jacobian()``. Instead, we only show two examples about the pure function and class function. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "253df55c",
   "metadata": {},
   "source": [
    "Given the following function, "
   ]
  },
  {
   "cell_type": "code",
   "id": "13ff570b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.355466Z",
     "start_time": "2025-10-06T03:32:48.351026Z"
    }
   },
   "source": [
    "import jax.numpy as jnp\n",
    "\n",
    "def f1(x, y):\n",
    "    a = 4 * x[1] ** 2 - 2 * x[2]\n",
    "    r = jnp.asarray([x[0] * y[0], 5 * x[2] * y[1], a, x[2] * jnp.sin(x[0])])\n",
    "    return r, a"
   ],
   "outputs": [],
   "execution_count": 41
  },
  {
   "cell_type": "code",
   "id": "1aefb47d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.898914Z",
     "start_time": "2025-10-06T03:32:48.373974Z"
    }
   },
   "source": [
    "_x = bm.array([1., 2., 3.])\n",
    "_y = bm.array([10., 5.])\n",
    "    \n",
    "grads, vec, aux = bm.jacobian(f1, return_value=True, has_aux=True)(_x, _y)"
   ],
   "outputs": [],
   "execution_count": 42
  },
  {
   "cell_type": "code",
   "id": "a6ea00cb",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.914496Z",
     "start_time": "2025-10-06T03:32:48.909924Z"
    }
   },
   "source": [
    "grads"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([[10.        ,  0.        ,  0.        ],\n",
       "                   [ 0.        ,  0.        , 25.        ],\n",
       "                   [ 0.        , 16.        , -2.        ],\n",
       "                   [ 1.6209068 ,  0.        ,  0.84147096]]),\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 43
  },
  {
   "cell_type": "code",
   "id": "c08984b8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.938058Z",
     "start_time": "2025-10-06T03:32:48.933272Z"
    }
   },
   "source": [
    "vec"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([10.       , 75.       , 10.       ,  2.5244129], dtype=float32)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 44
  },
  {
   "cell_type": "code",
   "id": "2b64116c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.961029Z",
     "start_time": "2025-10-06T03:32:48.955976Z"
    }
   },
   "source": [
    "aux"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(10., dtype=float32)"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 45
  },
  {
   "cell_type": "markdown",
   "id": "c1ad1eae",
   "metadata": {},
   "source": [
    "Given the following class objects,"
   ]
  },
  {
   "cell_type": "code",
   "id": "4f451a90",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:48.978918Z",
     "start_time": "2025-10-06T03:32:48.974211Z"
    }
   },
   "source": [
    "class Test(bp.BrainPyObject):\n",
    "  def __init__(self):\n",
    "    super().__init__()\n",
    "    self.x = bm.Variable(bm.array([1., 2., 3.]))\n",
    "\n",
    "  def __call__(self, y):\n",
    "    a = self.x[0] * y[0]\n",
    "    b = 5 * self.x[2] * y[1]\n",
    "    c = 4 * self.x[1] ** 2 - 2 * self.x[2]\n",
    "    d = self.x[2] * jnp.sin(self.x[0])\n",
    "    r = jnp.asarray([a, b, c, d])\n",
    "    return r, (c, d)"
   ],
   "outputs": [],
   "execution_count": 46
  },
  {
   "cell_type": "code",
   "id": "5f68ee77",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:49.078918Z",
     "start_time": "2025-10-06T03:32:48.993177Z"
    }
   },
   "source": [
    "t = Test()\n",
    "f_grad = bm.jacobian(t, grad_vars=t.x, argnums=0, has_aux=True, return_value=True)\n",
    "\n",
    "(var_grads, arg_grads), value, aux = f_grad(_y)"
   ],
   "outputs": [],
   "execution_count": 47
  },
  {
   "cell_type": "code",
   "id": "3db0d7d1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:49.089721Z",
     "start_time": "2025-10-06T03:32:49.084282Z"
    }
   },
   "source": [
    "var_grads"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[10.        ,  0.        ,  0.        ],\n",
       "       [ 0.        ,  0.        , 25.        ],\n",
       "       [ 0.        , 16.        , -2.        ],\n",
       "       [ 1.6209068 ,  0.        ,  0.84147096]], dtype=float32)"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 48
  },
  {
   "cell_type": "code",
   "id": "82547a2f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:49.110616Z",
     "start_time": "2025-10-06T03:32:49.105696Z"
    }
   },
   "source": [
    "arg_grads"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([[ 1.,  0.],\n",
       "                   [ 0., 15.],\n",
       "                   [ 0.,  0.],\n",
       "                   [ 0.,  0.]]),\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 49
  },
  {
   "cell_type": "code",
   "id": "382e1ab2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:49.132123Z",
     "start_time": "2025-10-06T03:32:49.127084Z"
    }
   },
   "source": [
    "value"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([10.       , 75.       , 10.       ,  2.5244129], dtype=float32)"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 50
  },
  {
   "cell_type": "code",
   "id": "de401f68",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T03:32:49.158820Z",
     "start_time": "2025-10-06T03:32:49.154360Z"
    }
   },
   "source": [
    "aux"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Array(10., dtype=float32), Array(2.5244129, dtype=float32))"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 51
  },
  {
   "cell_type": "markdown",
   "id": "8a486499",
   "metadata": {},
   "source": [
    "For more details on automatical differentation, please see our [API documentation](../apis/auto/math/autograd.rst)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "279.273px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
